from typing import List
from sympy import sympify
from latex2sympy2 import latex2sympy


class Formula:
    """公式计算"""
    def __init__(self, formula):
        self.formula = formula

    def render(self, variables: List[float]):
        symbols = list(self.formula.free_symbols)
        symbols.sort(key=lambda x: str(x))

        assert len(symbols) == len(variables), '参数和变量不匹配'
        value = self.formula.subs(zip(symbols, variables))
        return value


class SympyFormula(Formula):
    """公式计算
    eval 危险方法 可适配积分等公式
    """
    def __init__(self, formula_str):
        super(SympyFormula, self).__init__(formula=sympify(formula_str))


class Latex2Formula(Formula):
    """字符串公式转为sympy
    对于复杂公式转换能力有限， 尽量使用sympy创建公式
    """
    def __init__(self, formula_latex):
        formula = latex2sympy(formula_latex)
        super().__init__(formula=formula)


if __name__ == '__main__':
    sympy_formula = SympyFormula('integrate(exp(-x), (x, 0, oo))')  # 求积分公式
    result = sympy_formula.render([])
    print(result)

    latex_formula = Latex2Formula('a^3 + a * b')
    result = latex_formula.render([2, 2])
    assert result == 12, '计算结果错误'

    latex_formula = Latex2Formula('a^3 + a * b')
    result = latex_formula.render([2, 2])
    assert result == 12, '计算结果错误'

    latex_formula = Latex2Formula('a+b')
    result = latex_formula.render([1, 2])
    assert result == 3, '计算结果错误'

    latex_formula = Latex2Formula('a+b>1')
    result = latex_formula.render([1, 2])
    assert result == True, '计算结果错误'
